# -*- coding: UTF-8 -*-
import logging
import socket
import sys
import torch
import numpy as np
from PIL import Image
import io

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)


class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(1, 10, kernel_size=5),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2),
        )
        self.conv2 = torch.nn.Sequential(
            torch.nn.Conv2d(10, 20, kernel_size=5),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2),
        )
        self.fc = torch.nn.Sequential(
            torch.nn.Linear(320, 50),
            torch.nn.Linear(50, 10),
        )

    def forward(self, x):
        batch_size = x.size(0)
        x = self.conv1(x)  # 一层卷积层,一层池化层,一层激活层(图是先卷积后激活再池化，差别不大)
        x = self.conv2(x)  # 再来一次
        x = x.view(batch_size, -1)  # flatten 变成全连接网络需要的输入 (batch, 20,4,4) ==> (batch,320), -1 此处自动算出的是320
        x = self.fc(x)
        return x  # 最后输出的是维度为10的，也就是（对应数学符号的0~9）


def process_image(image_data):
    img = Image.open(io.BytesIO(image_data))
    img_array = np.array(img)
    model_input = torch.from_numpy(img_array).float().to(device)
    model = torch.load('./model_Mnist_optimal.pth', map_location=torch.device(device))
    model_input = model_input.unsqueeze(0)
    print(model_input.size())
    pre_res = model(model_input)
    pre_res = torch.argmax(pre_res, dim=1)
    return pre_res.item()


def handle_master_connection(master_socket):
    while True:
        try:
            # 接收图片数据长度
            data_length = master_socket.recv(1024)
            if not data_length:
                break
            data_length = int(data_length.decode())
            # 发送确认信息
            print("worker rec ", data_length)
            master_socket.sendall(b'OK')

            # 接收实际图片数据
            image_data = b''
            while len(image_data) < data_length:
                chunk = master_socket.recv(4096)
                if not chunk:
                    break
                image_data += chunk

            # 处理图片
            processed_res = process_image(image_data)
            print("处理结果：", processed_res)
            processed_res = str(processed_res).encode()
            # 发送处理结果
            master_socket.sendall(processed_res)

        except Exception as e:
            print(f"处理 Master 连接时出错: {e}")
            logging.exception(e)
            master_socket.sendall(b'ERROR')
            break


def register_with_master(master_host, master_port, worker_info):
    try:
        master_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        master_socket.connect((master_host, master_port))
        master_socket.sendall(f'REGISTER {worker_info}'.encode())
        response = master_socket.recv(1024)
        print(response.decode())
        if response == b'OK':
            print("已成功向 Master 注册")
            return master_socket
    except Exception as e:
        print(f"向 Master 注册出错: {e}")
    return None


def start_worker(master_host, master_port, worker_name):
    master_socket = register_with_master(master_host, master_port, worker_name)
    print(master_socket)
    if master_socket:
        handle_master_connection(master_socket)
        master_socket.close()
    else:
        print("启动出错")


if __name__ == "__main__":
    master_host = sys.argv[1]
    master_port = sys.argv[2]
    worker_name = socket.gethostname()
    start_worker(master_host, int(master_port), worker_name)
